# -*- encoding: utf-8 -*-
"""
        plot_olga_horizontal.py
        可视化功能函数
"""
import os
import matplotlib.pyplot as plt
import numpy as np
import matplotlib
import copy

matplotlib.rcParams["font.family"] = "SimHei"  # 或其他支持中文的字体
matplotlib.rcParams["axes.unicode_minus"] = False  # 解决负号'-'显示为方块的问题


def ploter(
    x: np.ndarray,
    y: np.ndarray,
    colors: list,
    label1: list,
    label2: list,
    linestyle: list,
    title: list,
    is_plot: bool,
is_dynamic_range:bool,
    savepath: str,
):
    """常规量可视化

    Args:
        x:np.ndarray;变量
        y:np.ndarray；变量
        colors:list;颜色
        label1:list;坐标轴命名
        label2:list;坐标轴命名
        linestyle:list;线类型
        title:list;名称
        savepath:str;保存路径

    Returns:
        None

    """
    # 常规量可视化
    t, nl, _ = y.shape
    plt.figure(figsize=(8, 5))
    
    for i_t in range(0, t):
        for i in range(0, nl):
            if i > len(linestyle) - 1:
                idx_line = -1
            else:
                idx_line = copy.deepcopy(i)
            plt.plot(
                x,
                y[i_t, i, ...],
                colors[i_t] + linestyle[idx_line],
                label=label1[i] + label2[i_t],
            )


    plt.xlim(0, np.max(x) * 1.1)
    if is_dynamic_range:
        plt.ylim(0, np.max(y) * 1.1)
   
    # 添加标签和标题
    plt.xlabel(title[0])
    plt.ylabel(title[1])
    plt.legend()

    if savepath:
        plt.savefig(savepath)
    # 显示图表
    if is_plot:
        plt.show()

    return

def ploterP(
    x: np.ndarray,
    y: np.ndarray,
    colors: list,
    label1: list,
    label2: list,
    linestyle: list,
    title: list,
    is_plot: bool,
is_dynamic_range:bool,
    savepath: str,
):
    """常规量可视化

    Args:
        x:np.ndarray;变量
        y:np.ndarray；变量
        colors:list;颜色
        label1:list;坐标轴命名
        label2:list;坐标轴命名
        linestyle:list;线类型
        title:list;名称
        savepath:str;保存路径

    Returns:
        None

    """
    # 常规量可视化
    t, nl, _ = y.shape
    plt.figure(figsize=(8, 5))
    
    for i_t in range(0, t):
        for i in range(0, nl):
            if i > len(linestyle) - 1:
                idx_line = -1
            else:
                idx_line = copy.deepcopy(i)
            plt.plot(
                x,
                y[i_t, i, ...],
                colors[i_t] + linestyle[idx_line],
                label=label1[i] + label2[i_t],
            )
    if is_dynamic_range:
        plt.xlim(0, np.max(x) * 1.1)
        plt.ylim(0, np.max(y) * 1.1)
   
    # 添加标签和标题
    plt.xlabel(title[0])
    plt.ylabel(title[1])
    plt.legend()

    if savepath:
        plt.savefig(savepath)
    # 显示图表
    if is_plot:
        plt.show()

    return

def ploterT(
    x: np.ndarray,
    y: np.ndarray,
    colors: list,
    label1: list,
    label2: list,
    linestyle: list,
    title: list,
    is_plot: bool,
is_dynamic_range : bool,
    savepath: str,
):
    """常规量可视化

    Args:
        x:np.ndarray;变量
        y:np.ndarray；变量
        colors:list;颜色
        label1:list;坐标轴命名
        label2:list;坐标轴命名
        linestyle:list;线类型
        title:list;名称
        savepath:str;保存路径

    Returns:
        None

    """
    # 常规量可视化
    t, nl, _ = y.shape
    plt.figure(figsize=(8, 5))
    
    for i_t in range(0, t):
        for i in range(0, nl):
            if i > len(linestyle) - 1:
                idx_line = -1
            else:
                idx_line = copy.deepcopy(i)
            plt.plot(
                x,
                y[i_t, i, ...],
                colors[i_t] + linestyle[idx_line],
                label=label1[i] + label2[i_t],
            )
    if is_dynamic_range:
        plt.xlim(0, np.max(x) * 1.1)
        plt.ylim(260, np.max(y) * 1.1)
   
    # 添加标签和标题
    plt.xlabel(title[0])
    plt.ylabel(title[1])
    plt.legend()

    if savepath:
        plt.savefig(savepath)
    # 显示图表
    if is_plot:
        plt.show()

    return

def vis_ploter(
    xall: np.ndarray,
    gt_flow_g,
    flow_g,
    gt_flow_L,
    flow_L,
    gt_p,
    p,
    gt_T,
    T,
    gt_alpha_L,
    alpha_L,
    times: list,
    mode: None,
    is_plot: bool,
is_dynamic_range:bool,
    savepath: None,
):
    """其他项

    Args:
        xall:np.ndarray X坐标
        ag:np.ndarray  气体截面分数计算值
        ul:np.ndarray  液体速度计算值
        ug:np.ndarray  气体速度计算值
        ag_as:np.ndarray 气体截面分数解析值
        ul_as:np.ndarray 液体速度解析值
        ug_as:np.ndarray 气体速度解析值
        times:list 时间点
        mode:None

    Returns:
        None
    """
    # 其他项
    data_flow_g = np.vstack([gt_flow_g, flow_g])
    data_flow_l = np.vstack([gt_flow_L, flow_L])
    data_p = np.vstack([gt_p, p])
    data_T = np.vstack([gt_T, T])
    data_alpha_L = np.vstack([gt_alpha_L, alpha_L])

    if not os.path.exists(savepath):
        os.makedirs(savepath)
    colors = ["r", "g", "b", "c", "m", "y", "k"]
    linestyle = ["-", "--", ":"]
    label1 = [
        "olga",
        "计算值",
    ]  # , '不可压', '可压']label1 = ['解析解', '不可压', '可压']
    label2 = [" t=" + str(t) for t in times]

    ploter(
        xall,
        np.asarray(data_flow_g)[None, ...],
        colors,
        label1,
        label2,
        linestyle,
        title=["网格", "气体流量 m3/s"], #"气体体积流量"
        is_plot=is_plot,
        is_dynamic_range = is_dynamic_range,
        savepath=savepath + "\气体流量.png",
    )
    ploterP(
        xall,
        np.asarray(data_p)[None, ...],
        colors,
        label1,
        label2,
        linestyle,
        title=["网格", "压力 kpa"],
        is_plot=is_plot,
        is_dynamic_range = is_dynamic_range,
        savepath=savepath + "\压力.png",
    )
    ploterT(
        xall,
        np.asarray(data_T)[None, ...],
        colors,
        label1,
        label2,
        linestyle,
        title=["网格", "温度 K"],
        is_plot=is_plot,
        is_dynamic_range = is_dynamic_range,
        savepath=savepath + "\温度.png",
    )
    ploter(
        xall,
        np.asarray(data_flow_l)[None, ...],
        colors,
        label1,
        label2,
        linestyle,
        title=["网格", "液体流量 m3/s"], #液体体积流量
        is_plot=is_plot,
        is_dynamic_range = is_dynamic_range,
        savepath=savepath + "\液体流量.png",
    )
    ploter(
        xall,
        np.asarray(data_alpha_L)[None, ...],
        colors,
        label1,
        label2,
        linestyle,
        title=["网格", "持液率 0~1"],
        is_plot=is_plot,
        is_dynamic_range = is_dynamic_range,
        savepath=savepath + "\持液率.png",
    )
    return

# 颜色代码
text_colors = {
    "black": 30,
    "red": 31,
    "green": 32,
    "yellow": 33,
    "blue": 34,
    "magenta": 35,
    "cyan": 36,
    "white": 37,
}

bg_colors = {
    "black": 40,
    "red": 41,
    "green": 42,
    "yellow": 43,
    "blue": 44,
    "magenta": 45,
    "cyan": 46,
    "white": 47,
}


def colored_text(text, text_color, bg_color):
    """为打印加上底色

    Args:
        text (_type_): _description_
        text_color (_type_): _description_
        bg_color (_type_): _description_

    Returns:
        _type_: _description_
    """
    return f"\033[{text_color};{bg_color}m{text}\033[0m"
